# -*- coding: UTF-8 -*-
import os
import netCDF4 as nc
import numpy as np
from osgeo import gdal, osr, ogr
import glob
from tqdm import tqdm


# os.environ['PROJ_LIB'] = r'./proj'

# 用于逐月数据集

def nc2tifs(nc_file: str, output_folder: str, i: int, key: str, scale=0.1):
    """
    :param nc_file: 多波段/多维度nc文件完整路径
    :param output_folder: 输出文件夹
    :param i: 第i个维度或波段
    :param key: nc文件的索引
    :param scale：比例系数，单位缩放，默认是0.1
    :return: 保存的tif文件名
    """
    data = nc.Dataset(nc_file)  # 利用Dataset()读取nc数据
    Lat = data.variables['lat'][:]
    Lon = data.variables['lon'][:]
    nodata =data.variables[key].missing_value  # netcdf文件的缺失值
    data_arr = np.asarray(data.variables[key])  # 属性变量名
    mask = np.where(data_arr[0, :, :] == nodata, 0, 1)  # nodata区域为 0, 多波段只读取一个波段
    _mask = np.where(mask == 0, 1, 0)  # mask的反
    data_arr = data_arr * mask * scale + _mask * nodata   # 把nodata数据添加上
    # 影像的左上角&右下角坐标
    Lonmin, Latmax, Lonmax, Latmin = [Lon.min(), Lat.max(), Lon.max(), Lat.min()]
    # Lonmin, Latmax, Lonmax, Latmin
    # 分辨率计算
    Num_lat = len(Lat)
    Num_lon = len(Lon)
    Lat_res = (Latmax - Latmin) / (float(Num_lat) - 1)
    Lon_res = (Lonmax - Lonmin) / (float(Num_lon) - 1)
    # 创建tif文件
    driver = gdal.GetDriverByName('GTiff')
    out_tif_name = output_folder + '\\' + nc_file.split('\\')[-1].split('.')[0]  + format(i + 1, '02') + '.tif'
    out_tif = driver.Create(out_tif_name, Num_lon, Num_lat, 1, gdal.GDT_Float32)
    # 设置影像的显示范围
    # Lat_re前需要添加负号
    geotransform = (Lonmin, Lon_res, 0.0, Latmax, 0.0, -Lat_res)
    out_tif.SetGeoTransform(geotransform)
    # 定义投影
    prj = osr.SpatialReference()
    prj.ImportFromEPSG(4326)
    out_tif.SetProjection(prj.ExportToWkt())
    out_tif.GetRasterBand(1).SetNoDataValue(nodata)
    # 数据导出
    out_tif.GetRasterBand(1).WriteArray(data_arr[i, :, :])  # 将数据写入内存
    out_tif.FlushCache()  # 将数据写入到硬盘
    out_tif = None  # 关闭tif文件

    return out_tif_name


if __name__ == "__main__":

    # #测试数据信息，每一个nc文件是逐月降水、均温数据
    # #中国1km分辨率逐月降水量数据集（1901 - 2023）
    # #https://data.tpdc.ac.cn/zh-hans/data/faae7605-a0f2-4d18-b28f-5cee413766a2
    #中国1km分辨率逐月平均气温数据集（1901-2023）
    #https://data.tpdc.ac.cn/zh-hans/data/71ab4677-b66c-4fd1-a004-b2a541c4d5bf

    Input_folder = r"D:\数据处理\TMP"
    Output_folder = r"D:\数据处理\TMP2"
    key = 'tmp'  # 均温是tmp,pre,etp
    time = 12  # 时间维度，每一个
    scale = 0.1  # 降水比例转换0.1mm，温度比例转换 0.1deg，pet单位0.1mm

    if os.path.exists(Output_folder) == False:  # 如果目录不存在就返回False，创建该文件夹
        os.mkdir(Output_folder)

        # 读取所有数据
    nc_list = sorted(glob.glob(os.path.join(Input_folder, f"*.nc")))
    # print(nc_list)
    for nc_file in tqdm(nc_list):  # 循环每一个文件
        # data = nc.Dataset(nc_file)  # 利用.Dataset()读取nc数据
        for i in range(12):
            out_tif_name = nc2tifs(nc_file, Output_folder, i, key)
            # print(out_tif_name + '\n')

    # #代码
    # data = nc.Dataset(r'D:\Users\Jiudian\Desktop\pre\pre_2000.nc', 'r', format = 'NETCDF4')
    # print(data.variables.keys()) # 查看变量
    # # 输出 dict_keys(['lon', 'lat', 'time', 'pre'])
    #
    # print(data.variables['pre'].unit) # 查看降水的单位
    # # 输出 0.1 mm
    #
    # print(data.variables['pre'][:].shape) # 查看数据维度
    # # 输出 (12, 5146, 7849)
    #
    # print(data.variables['time'][:].shape) # 查看数据维度
    # # 输出 (12, )
    # print(data.variables['time'][:]) # 查看数据维度
    #
    # # 输出 [ 1.  2.  3.  4.  5.  6.  7.  8.  9. 10. 11. 12.]
    # print(int(data.variables['time'][0])) # 查看数据维度